import numpy as np
import mindspore.nn as nn
import mindspore.ops.operations as P
from mindspore import Tensor
from mindspore.nn.cell import Cell
from mindspore.common.initializer import Normal
#from src.dataset1 import MovieReview


class DPCNN(nn.Cell):
    """Lenet network structure."""
    # define the operator required
    def __init__(self, num_class=10, num_channel=1):
        super(DPCNN, self).__init__()
        self.unsqueeze = P.ExpandDims()
        self.embedding = nn.Embedding(18885, 40)

        self.conv_region=nn.Conv2d(1, 250, (3,40),pad_mode='valid',weight_init='normal', has_bias=True)
        self.conv = nn.Conv2d(250, 250, (3,1), pad_mode='valid',weight_init='normal', has_bias=True)
        self.conv2 = nn.Conv2d(250, 250, (3,1), pad_mode='same')
        self.max_pool2d = nn.MaxPool2d(kernel_size=(3, 1), stride=2)
        self.max_pool2d2 = nn.MaxPool2d(kernel_size=(49, 1))

        self.relu = nn.ReLU()
        self.fc1 = nn.Dense(250, 2, weight_init=Normal(0.02))
        self.fc2 = nn.Dense(96, 2, weight_init=Normal(0.02))
        self.fc3 = nn.Dense(60, 2, weight_init=Normal(0.02))
        self.flatten = nn.Flatten()

        self.padding1 = nn.Pad(paddings=((0,0),(0,0),(1,1),(0,0)))
        self.padding2 = nn.Pad(paddings=((0,0),(0,0),(0,1),(0,0)))
        self.drop = nn.Dropout(keep_prob=0.5)
        self.reducemean = P.ReduceMax(keep_dims=False)

    def construct(self, x):
        #print('ka')
        x = self.unsqueeze(x, 1)
        x = self.embedding(x)
        #print('embedding', x.shape)
        # 1卷积+2same卷积
        x = self.conv_region(x)
        #print('conv_region', x.shape)
        x=self.padding1(x)
        #print('padding1', x.shape)
        x = self.relu(x)

        x = self.conv(x)
        #print('conv1', x.shape)
        x = self.padding1(x)
        #print('padding2', x.shape)
        x = self.relu(x)
        x = self.conv(x)
        #print('conv2', x.shape)
        while x.shape[2] > 2:
            x = self._block(x)

        #print(1)
        x = self.flatten(x)
        x = self.drop(x)

        x = self.fc1(x)
        #print('finall',x.shape)
        return x

    def _block(self, x):

        x=self.padding2(x)
        px = self.max_pool2d(x)

        x = self.padding1(px)
        x = self.relu(x)
        x = self.conv(x)

        x = self.padding1(x)
        x = self.relu(x)
        x = self.conv(x)
        x=x+px
        #print('_block', x.shape)

        # px = self.max_pool2d(x)
        # print('px',px.shape)
        # x = self.relu(px)
        # x = self.conv2(x)
        # print('conv2', x.shape)
        # x = self.relu(x)
        # x = self.conv2(x)
        # print('conv2', x.shape)
        # x = x + px
        # print('_block',x.shape)
        return x

# images=Tensor(images)
# LeNet5().construct(images)